from d2l import torch as d2l
import torch

#返回预测正确的个数
def accuracy(y_hat,y):
    if len(y_hat.shape)>1 and y_hat.shape[1]>1:
        y_hat = y_hat.argmax(axis=1)
    cmp = y_hat.type(y.dtype) == y #统一数据格式
    true_num = cmp.sum()
    return float(true_num)

def evaluate_accuracy(net,data_iter):
    if isinstance(net,torch.nn.Module):
        net.eval() #将模型设置为评估模式
    metric = d2l.Accumulator(2) #正确预测数，样本总数
    for X,y in data_iter:
        if isinstance(net,torch.nn.Module):
            y_hat = net(X)
        else:
            y_hat = net(X,False)
        metric.add(accuracy(y_hat,y),y.numel())
    return metric[0] / metric[1]

def train_epoch(net,train_iter,loss_func,trainer):
    if isinstance(net,torch.nn.Module):
        net.train()
    metric = d2l.Accumulator(3)
    for X,y in train_iter:
        if isinstance(net,torch.nn.Module):
            y_hat = net(X)
        else:
            y_hat = net(X,True)
        l = loss_func(y_hat,y)
        #l：样本平均误差
        if isinstance(trainer,torch.optim.Optimizer):
            trainer.zero_grad()
            l.backward()
            trainer.step()
            metric.add(float(l)*len(y),accuracy(y_hat,y),y.numel())
        else:
            l.sum().backward()
            trainer()
            metric.add(float(l.sum()),accuracy(y_hat,y),y.numel())
    return metric[0]/metric[2],metric[1]/metric[2]

def train(net,train_iter,test_iter,loss_func,num_epochs,trainer):
    animator = d2l.Animator(xlabel='epoch',xlim=[1,num_epochs],ylim=[0.3,0.9],
                            legend=['train_loss','train_accuracy','test_accuracy'])
    for epoch in range(num_epochs):
        train_metrics = train_epoch(net,train_iter,loss_func,trainer)
        test_acc = evaluate_accuracy(net,test_iter)
        animator.add(epoch,train_metrics+(test_acc,))
        print('epoch:',epoch,'loss:',train_metrics[0])
    train_loss,train_acc = train_metrics
    d2l.plt.show()

def predict(net,test_iter,n=6):
    for X,y in test_iter:
        break
    trues = d2l.get_fashion_mnist_labels(y)
    if isinstance(net,torch.nn.Module):
        y_hat = net(X)
    else:
        y_hat = net(X,False)
    predicts = d2l.get_fashion_mnist_labels(y_hat.argmax(axis=1))
    titles = [true+'\n'+predict for true,predict in zip(trues,predicts)]
    d2l.show_images(X[:n].reshape((n,28,28)),1,n,titles=titles[0:n])
    d2l.plt.show()